/* By Ventrue  */
#include <linux/sched.h>
#include <linux/rbtree.h>
#include "sched.h"

/* 更新截止时间 */
static void update_deadline(struct rms_entity *rms_e) {
    rms_e->deadline += rms_e->period;
}

/* 比较两个任务的优先级(周期小prio大)返回布尔值 a是否优先级>b*/
static bool rms_task_prio_cmp(struct rb_node *a, const struct rb_node *b)
{
    struct rms_entity *ent_a = rb_entry(a, struct rms_entity, run_node);
    const struct rms_entity *ent_b = rb_entry(b, struct rms_entity, run_node);
    return ent_a->period < ent_b->period;
}

/* 将任务加入就绪队列 */
static void rms_enqueue_task(struct rq *rq, struct task_struct *p, int flags) {
    struct rms_entity *rms_e = &p->rms;
    struct rms_rq *rms = &rq->rms;
    
    raw_spin_lock(&rms->lock);
    update_deadline(rms_e); // 更新截止时间为下一个周期
    rb_add(&rms_e->run_node, &rms->ready_queue, rms_task_prio_cmp);
    raw_spin_unlock(&rms->lock);
}

/* 从队列中移除任务 */
static void rms_dequeue_task(struct rq *rq, struct task_struct *p, int flags) {
    struct rms_entity *rms_e = &p->rms;
    struct rms_rq *rms = &rq->rms;
    
    raw_spin_lock(&rms->lock);
    rb_erase(&rms_e->run_node, &rms->ready_queue);
    RB_CLEAR_NODE(&rms_e->run_node);
    raw_spin_unlock(&rms->lock);
}

/* 选择下一个要执行的任务 */
static struct task_struct *rms_pick_next_task(struct rq *rq) {
    struct rms_rq *rms = &rq->rms;
    struct rb_node *left = rb_first(&rms->ready_queue);

    if (node) {
        struct rms_entity *e = rb_entry(node, struct rms_entity, run_node);
        printk(KERN_INFO "RMS: Picked task %s (period=%llu)\n",
               container_of(e, struct task_struct, rms)->comm, e->period);
    }

    if (!left)
        return NULL;
    
    struct rms_entity *rms_e = rb_entry(left, struct rms_entity, run_node);
    return container_of(rms_e, struct task_struct, rms);
}

/* 时间片到期处理 */
static void rms_task_tick(struct rq *rq, struct task_struct *p, int queued) {
    struct rms_entity *rms_e = &p->rms;
    
    // 检查是否超过截止时间
    if (unlikely(rq->clock > rms_e->deadline)) {
        printk(KERN_WARNING "RMS: Task %s missed deadline!\n", p->comm);
        // 此处可添加错误处理逻辑
    }
}

/* 检查是否抢占当前任务 */
static void rms_check_preempt_curr(struct rq *rq, struct task_struct *p, int flags)
{
    if (p->sched_class == &rms_sched_class) {
        struct task_struct *curr = rq->curr;
        if (curr->sched_class != &rms_sched_class)
            resched_curr(rq);  // RMS任务抢占非RMS任务
        else if (rms_task_prio_cmp(&p->rms.run_node, &curr->rms.run_node))
            resched_curr(rq);  // 周期更短的RMS任务抢占当前任务
    }
}

/* 设置下一个任务 */
static void rms_set_next_task(struct rq *rq, struct task_struct *next)
{
    struct rms_entity *rms_e = &next->rms;
    rms_e->deadline = ktime_get_ns() + rms_e->period;  // 重置截止时间
    update_curr_rms(rq);  // 更新统计信息（需实现）
}

/* 初始化RMS运行队列 */
void init_rms_rq(struct rms_rq *rms_rq) {
    rms_rq->ready_queue = RB_ROOT;
    raw_spin_lock_init(&rms_rq->lock);
    rms_rq->nr_running = 0;
}

__init void init_sched_rms_class(void)
{

}

/* 定义RMS调度类 */
/* 按照段的排序 并声明了rms调度算法具体的段名称 */
const struct sched_class rms_sched_class 
__section("__rms_sched_class")= {
    .enqueue_task   = rms_enqueue_task,
    .dequeue_task   = rms_dequeue_task,
    .pick_next_task = rms_pick_next_task,
    .task_tick      = rms_task_tick,
    .check_preempt_curr = rms_check_preempt_curr,
    .set_next_task = rms_set_next_task,

};
